from typing import Dict, List, Set
import os

class TreeVisualizer:
    """Visualizes function call trees and other analyses as text."""
    
    @staticmethod
    def visualize_call_tree(tree, output_file):
        """
        Visualize a function call tree and save it to a file.
        
        Args:
            tree: Call tree dictionary
            output_file: Path to save the visualization
        """
        with open(output_file, 'w') as f:
            for root_func, subtree in tree.items():
                # Write the root function name
                f.write(f"{root_func}\n")
                
                # Handle special case when there's a main subfunction at the top level
                if "main" in subtree:
                    # Skip directly to main's children, bypassing main itself
                    main_subtree = subtree["main"]
                    # Don't include the actual main function in the tree
                    for sub_key in sorted(main_subtree.keys()):
                        # Skip recursive main references
                        if sub_key == "main" or "(recursive)" in sub_key:
                            continue
                        # For each child of main, print it as a direct child of the root
                        f.write(f"├── {sub_key}\n")
                        # Process its subtree
                        if main_subtree[sub_key]:
                            TreeVisualizer._print_subtree(main_subtree[sub_key], f, "│   ", "")
                    
                    # Process other children of root (not main)
                    other_keys = [k for k in sorted(subtree.keys()) if k != "main"]
                    for i, key in enumerate(other_keys):
                        is_last = (i == len(other_keys) - 1)
                        branch = "└── " if is_last else "├── "
                        f.write(f"{branch}{key}\n")
                        if subtree[key]:
                            new_prefix = "    " if is_last else "│   "
                            TreeVisualizer._print_subtree(subtree[key], f, new_prefix, branch)
                else:
                    # Normal processing for other root functions
                    TreeVisualizer._print_subtree(subtree, f, "", "")
        
        print(f"Call tree saved to {output_file}")
    
    @staticmethod
    def _print_subtree(tree, file, prefix, branch_prefix):
        """Recursively print a subtree."""
        if not tree:
            return
        
        # Sort keys for consistent output
        sorted_keys = sorted(tree.keys())
        
        # Filter out push_back nodes and main function when it appears as a subfunction
        sorted_keys = [key for key in sorted_keys if key != "push_back" and key != "main"]
        
        for i, key in enumerate(sorted_keys):
            # Check if this is the last item
            is_last = (i == len(sorted_keys) - 1)
            
            # Choose the appropriate branch character
            branch = "└── " if is_last else "├── "
            
            # Check if this is a recursive call to a function already in the path
            if "(recursive)" in key:
                file.write(f"{prefix}{branch}{key}\n")
                continue
                
            # Check for repeat nodes to avoid duplication
            if "(repeat)" in key:
                file.write(f"{prefix}{branch}{key}\n")
                continue
            
            file.write(f"{prefix}{branch}{key}\n")
            
            # Recursively print the subtree
            if tree[key]:
                # Continue the branch character pattern for children
                new_prefix = prefix + ("    " if is_last else "│   ")
                TreeVisualizer._print_subtree(tree[key], file, new_prefix, branch)
    
    @staticmethod
    def visualize_computation_steps(steps, output_file):
        """
        Visualize computation steps and save to a file.
        
        Args:
            steps: Dictionary of computation steps for each function
            output_file: Path to save the visualization
        """
        with open(output_file, 'w') as f:
            for func_name, func_steps in steps.items():
                f.write(f"Function: {func_name}\n")
                f.write("=" * (len(func_name) + 10) + "\n\n")
                
                for i, step in enumerate(func_steps, 1):
                    f.write(f"Step {i}: {step}\n")
                
                f.write("\n" + "=" * 60 + "\n\n")
        
        print(f"Computation steps saved to {output_file}")
    
    @staticmethod
    def visualize_call_chains(chains, output_file):
        """
        Visualize call chains and save to a file.
        
        Args:
            chains: List of call chains
            output_file: Path to save the visualization
        """
        with open(output_file, 'w') as f:
            for i, chain in enumerate(chains, 1):
                f.write(f"Chain {i}: {' -> '.join(chain)}\n")
        
        print(f"Call chains saved to {output_file}") 